from vanna.chromadb import ChromaDB_VectorStore
from vanna.flask import VannaFlaskApp
from guijiAI_chat import GuijiAI_Chat

class MyVanna(ChromaDB_VectorStore, GuijiAI_Chat):
    def __init__(self, config=None):
        ChromaDB_VectorStore.__init__(self, config=config)
        GuijiAI_Chat.__init__(self, config=config)




def main():
    print("Hello from vannai!")
    vn = MyVanna(config={'api_key':'sk-qwldskhmatxlujmhqkdqhvfriiijewwxtfsxecocgbfqtfia','model': 'Qwen/Qwen2.5-7B-Instruct'})
#     vn.train(
#         ddl="""
#         CREATE TABLE transaction_records (
#     feedback_unit VARCHAR(100) COMMENT '反馈单位',
#     transaction_time TIMESTAMP(6) COMMENT '交易时间',
#     name VARCHAR(100) COMMENT '名称',
#     our_account VARCHAR(50) COMMENT '本方账号',
#     our_card_number VARCHAR(50) COMMENT '本方卡号',
#     id_number VARCHAR(50) COMMENT '证件号码',
#     transaction_type VARCHAR(50) COMMENT '交易类型',
#     debit_credit_flag VARCHAR(10) COMMENT '借贷标志',
#     currency VARCHAR(10) COMMENT '币种',
#     transaction_amount NUMERIC(15, 2) COMMENT '交易金额',
#     flow_direction VARCHAR(20) COMMENT '流向',
#     transaction_balance NUMERIC(15, 2) COMMENT '交易余额',
#     counterparty_name VARCHAR(100) COMMENT '交易对方名称',
#     counterparty_account VARCHAR(50) COMMENT '交易对方账号',
#     counterparty_card_number VARCHAR(50) COMMENT '交易对方卡号',
#     counterparty_id_number VARCHAR(50) COMMENT '交易对方证件号码',
#     counterparty_bank VARCHAR(100) COMMENT '交易对方账号开户行',
#     transaction_summary VARCHAR(200) COMMENT '交易摘要',
#     branch_name VARCHAR(100) COMMENT '交易网点名称',
#     transaction_location VARCHAR(100) COMMENT '交易发生地',
#     merchant_name VARCHAR(100) COMMENT '商户名称',
#     ip_address VARCHAR(50) COMMENT 'IP地址',
#     mac_address VARCHAR(50) COMMENT 'MAC地址',
#     teller_number VARCHAR(50) COMMENT '交易柜员号',
#     remarks TEXT COMMENT '备注'
# );
#         """
#     )
    # vn.connect_to_(host='127.0.0.1', dbname='database', user='postgres', password='passmord', port-5432)
    app = VannaFlaskApp(vn)
    app.run()
    return app


if __name__ == "__main__":
    main()
